#ifndef ANAGRAM_OPTIMIZER_NGRAMTOOL_NGRAM_TABLE_H
#define ANAGRAM_OPTIMIZER_NGRAMTOOL_NGRAM_TABLE_H

#include <cstdio>
#include <cstring>
#include <unordered_map>

#include "ngram.h"
#include "istring.h"

namespace anagram_optimizer {

template<int N>
class NGramTable {
 public:
  // The following options may be ORed together and passed to Read().
  static const int kSumFrequencies = 0x0001;  // Default: error on dupes.

  NGramTable() {}

  void Add(const NGram<N>& ngram, int count) {
    frequencies_[ngram] += count;
  }

  int Count(const NGram<N>& ngram) {
    return frequencies_[ngram];
  }

  bool Write(const std::string& filename);
  bool Read(const std::string& filename, unsigned int options=0);

 private:
  // Disallow copy and assign.
  NGramTable(const NGramTable& other) {}
  NGramTable& operator =(const NGramTable& other) {}

  std::unordered_map<NGram<N>, size_t> frequencies_;
};

template<int N>
bool NGramTable<N>::Read(const std::string& filename, unsigned int options) {
  frequencies_.clear();

  FILE* fp = fopen(filename.c_str(), "r");
  if (fp == NULL) {
    perror("fopen: ");
    return false;
  }

  static const char kWhitespace[] = " \t\r\n";

  char buf[512];
  std::string line;
  int line_number = 0;
  while (!feof(fp)) {
    ++line_number;

    // Read a whole line at a time out of the file.
    //
    // TODO(drk):  I am a bastard. 
    line.clear();
    do {
      buf[0] = '\0';
      buf[511] = '\n';
      if (!fgets(buf, sizeof(buf), fp)) {
        // End of file (or error, which we silently ignore ... heh).
        line += buf;
        break;
      }
      line += buf;
      // Now, buf[511] == '\n' -> whole line was read, but not the converse,
      // since buf[511] == '\0' also -> whole line was read.  Together, these
      // completely account for possible scenarios where the whole line was
      // read.
    } while (buf[511] != '\n' && buf[511] != '\0');
    if (line.empty()) break;

    // Each line consists of the words comprising the n-gram (separated by one
    // or more whitespace characters) and a frequency (a non-negative
    // integer).
    NGram<N> ngram;
    size_t start = line.find_first_not_of(kWhitespace);;
    size_t end = line.find_first_of(kWhitespace, start);
    int word_count = 0;
    while (start != std::string::npos) {
      if (word_count == N + 1) {
        fprintf(stderr, "%s:%d: too many elements\n", filename.c_str(),
                line_number);
        fclose(fp);
        return false;
      }
      if (end <= start) {
        // Empty line.
        break;
      }

      std::string word(line, start, end - start);
      ++word_count;
      size_t frequency;
      if (word_count == N + 1) {
        // Parse the frequency, with special care to make sure it's a valid,
        // non-negative number.
        //
        // TODO(drk):  This is kind of squirrelly.  Plus, strtoul actually
        // silently converts negative numbers to their unsigned long int
        // equivalent value, which sucks.
        char* buf = new char[word.size() + 1];
        char* endptr = NULL;
        strcpy(buf, word.c_str());
        frequency = strtoul(buf, &endptr, 10);
        delete [] buf;
        if (endptr != NULL &&
            static_cast<size_t>(endptr - buf) < word.size()) {
          fprintf(stderr, "%s:%d: cannot parse frequency \"%s\"\n",
                  filename.c_str(), line_number, word.c_str());
          fclose(fp);
          return false;
        }
        if (word[0] == '-') {
          fprintf(stderr, "%s:%d: negative frequency is invalid\n",
                  filename.c_str(), line_number);
          fclose(fp);
          return false;
        }

        // Handling of repeated entries for the same n-gram is conditional on
        // an option flag.  Normally we error out (after all, it seems odd to
        // have a dupe in a curated data file), but optionally we just sum
        // frequencies together.
        auto it = frequencies_.find(ngram);
        if (it != frequencies_.end()) {
          if (options & kSumFrequencies) {
            it->second += frequency;
          } else {
            fprintf(stderr, "%s:%d: repeated ngram\n",
                    filename.c_str(), line_number);
            fclose(fp);
            return false;
          }
        } else {
          frequencies_[ngram] = frequency;
        }
      } else {
        ngram.words[word_count - 1] = istring(word);
      }

      // Finally, advance to the next word.
      start = line.find_first_not_of(kWhitespace, end);
      end = line.find_first_of(kWhitespace, start);
    }
    if (word_count > 0 && word_count != N + 1) {
      fprintf(stderr, "%s:%d: missing elements\n",
              filename.c_str(), line_number);
      fclose(fp);
      return false;
    }
  }
  fclose(fp);
  return true;
}

template<int N>
bool NGramTable<N>::Write(const std::string& filename) {
  FILE* fp = fopen(filename.c_str(), "w");
  if (fp == NULL) {
    perror("fopen: ");
    return false;
  }

  for (const auto& p : frequencies_) {
    for (int i = 0; i < N; ++i) {
      if (fputs(p.first.words[i].as_string().c_str(), fp) == EOF ||
          fputc(' ', fp) == EOF) {
        fprintf(stderr, "Error writing to file \"%s\".", filename.c_str());
        fclose(fp);
        return false;
      }
    }
    if (fprintf(fp, "%zd\n", p.second) < 0) {
      fprintf(stderr, "Error writing to file \"%s\".", filename.c_str());
      fclose(fp);
      return false;
    }
  }
  fclose(fp);
  return true;
}
}  // anagram_optimizer

#endif // ANAGRAM_OPTIMIZER_NGRAMTOOL_NGRAM_TABLE_H
